/**
nschaind.c
Copyright Christiaan Ottow 2009
chris@6core.net

See README for info
 
 This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <stdlib.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <stdint.h>
#include <string.h>
#include <ldns/ldns.h>

#define PORT 53
#define TTL (3600*24*7)

#define CHAIN_COUNT 10

#define TRUE 1
#define FALSE 0

#define BUFLEN 10000

bool have_victim = 0;
in_addr_t victim;
int cname_counter = 1;
int chain_counter = 1;
long last_time = 0;

uint16_t predict1[10];
uint16_t predict2[10];
bool prediction = 0;
int prediction_count = 0;
short victim_ttl = 3;
int victim_sports = 0;
uint16_t last_sport;

/* variables from commandline */
char *cname_base;
char *mxname;
char *final_record;

void diep(char *s)
{
	perror(s);
	exit(1);
}

/**
 * Where the magic happens
 */
void predict_ids( int bind_version, uint16_t qid, uint16_t *cand)
{
	// return values
	//uint16_t *cand = malloc( 10 * sizeof( uint16_t ) );

	// shift registers
	int tap1 = 0x80000057;
	int tap2 = 0x80000062;
	
	if( bind_version == 0 )
	{
		tap1 = ( tap1 >> 1 ) | ( 1 << 31 );
		tap2 = ( 0x800000c2 >> 1 ) | ( 1 << 31 );
	}
	
	// one bit shift (assuming the two lsb's are 0 and 0 )
	int msb,bufp=0;
	for( msb = 0; msb < 2; msb++ )
	{
		//printf("SIT1: setting index %d\n", bufp );
		*(cand+bufp) = ( ( msb << 15 ) | (qid >> 1 ) ) & 0xffff;
		bufp++;
	}
	
	// two bit shift (assuming the two lsb's are 1 and 1 )
	// first shift (we know the lsb is 1 in both LFSRs)
	uint16_t v = qid, v1, v2;
	v =  ( (v >> 1 ) ^ tap1 ) ^ tap2;
	if( (v & 1 ) == 0 )
	{
		// after the first shift, the lsb becomes 0, so the two LFSRs now have identical lsbs: 0 and 0 or 1 and 1
		// second shift:
		v1 = (v >> 1 );				// 0 and 0
		v2 = ( (v >> 1 ) & tap1 ) ^ tap2;	// 1 and 1
	} else {
		// after the first shift, the lsb becomes 1, so the two LFSRs now have different lsb's: 1 and 0 or 0 and 1
		// second shift:
		v1 = (v >> 1 ) ^ tap1;
		v2 = (v >> 1 ) ^ tap2;
	}
	
	// also need to enumerate over the 2 msb's we are clueless about
	int msbits;
	for( msbits=0; msbits < 4; msbits++)
	{
		//printf("SIT2: settings index %d\n", bufp );
		*(cand+bufp) = ( ( msbits << 14 ) | v1 ) & 0xffff;
		bufp++;
		//printf("SIT3: settings index %d\n", bufp );
		*(cand+bufp) = ( ( msbits << 14 ) | v2 ) & 0xffff;
		bufp++;
	}
}

/**
 *  Returns a RR which answers the question with a CNAME to the next hostname
 */
ldns_rr *a_record( bool do_chain, char *questionstr )
{
	ldns_rr *answer = ldns_rr_new();
	char answerstr[2000];
	if( do_chain == TRUE )
	{
		sprintf( answerstr, "%s %d IN CNAME chain%02d.%s", questionstr, TTL, cname_counter, cname_base);
	} else {
		sprintf( answerstr, "%s %d IN CNAME %s", questionstr, TTL, final_record );
	}
	
	cname_counter++;
	ldns_rdf *prev = NULL;
	int status = ldns_rr_new_frm_str( &answer, answerstr, TTL, NULL, &prev );
	if (status != LDNS_STATUS_OK ) {
		printf( "Error while creating rr: %s\n", ldns_get_errorstr_by_id(status));
	}
	return answer;
}

/**
 * Returns MXNAME as the mailserver and gives it a CNAME
 */
ldns_rr_list *mx_answer( bool do_chain, ldns_pkt *p, char *questionstr )
{
	ldns_rr *mx_answer = ldns_rr_new();
	ldns_rr_set_class( mx_answer, LDNS_RR_CLASS_IN );
	ldns_rr_set_type( mx_answer, LDNS_RR_TYPE_MX );
	char mxstr[2000];
	sprintf( mxstr, "%s %d IN MX 10 %s.%s", questionstr, TTL, mxname, cname_base );
	ldns_rdf *prev = NULL;
	int status = ldns_rr_new_frm_str( &mx_answer, mxstr, TTL, NULL, &prev );
	if( status != LDNS_STATUS_OK ) {
		printf( "Error while creating rr: %s\n", ldns_get_errorstr_by_id( status ));
	}

	char fullmx[100] = "";
	strcat( fullmx, mxname );
	strcat( fullmx, ".");
	strcat( fullmx, cname_base );
	
	ldns_rr *a_answer = a_record( do_chain, fullmx );

	ldns_rr_list *answerlist = ldns_rr_list_new();
	ldns_rr_list_push_rr( answerlist, mx_answer );
	ldns_rr_list_push_rr( answerlist, a_answer );
	ldns_pkt_set_ancount( p, 2 );
	
	return answerlist;
}

ldns_rr_list *a_answer( bool do_chain, ldns_pkt *p, char *questionstr )
{	
	ldns_rr *answer = a_record( do_chain, questionstr );
	
	ldns_rr_list *answerlist = ldns_rr_list_new();
	ldns_rr_set_class( answer, LDNS_RR_CLASS_IN );
	ldns_rr_set_type( answer, LDNS_RR_TYPE_CNAME );
	ldns_rr_list_push_rr( answerlist, answer );
	ldns_pkt_set_ancount( p, 1 );

	return answerlist;
}

ldns_pkt *create_response_pkt( bool do_chain, ldns_pkt *orig )
{
	// get request data from request
	ldns_rr_list *questionlist = ldns_pkt_question( orig );
	ldns_rr *question = ldns_rr_list_pop_rr( questionlist );
	ldns_rdf *dname = ldns_rr_owner( question );
	ldns_rr_list_push_rr( questionlist, question );
	char *questionstr = ldns_rdf2str( dname );
	int qtype = ldns_rr_get_type( question );
	
	ldns_pkt *p = ldns_pkt_new();
	
	ldns_rr_list *answerlist;
	
	// set id
	ldns_pkt_set_id( p, ldns_pkt_id( orig ) );
	
	// set header fields
	// id qr opcode aa tc rd ra z rcode qdcount ancount nscount arcount
	ldns_pkt_set_qr( p, 1 );							// question or answer
	ldns_pkt_set_opcode( p, LDNS_PACKET_QUERY );		// type of response
	ldns_pkt_set_aa( p, 1 );							// authorative answer
	ldns_pkt_set_tc( p, 0 );							// truncated
	ldns_pkt_set_rd( p, 0 );							// recursion desired
	ldns_pkt_set_ra( p, 0 );							// recursion available
	ldns_pkt_set_rcode( p, 0);							// error condition
	ldns_pkt_set_qdcount( p, 1 );						// question count
	ldns_pkt_set_nscount( p, 0 );						// nameserver count
	ldns_pkt_set_arcount( p, 0 );						// authority count
	
	if( qtype == LDNS_RR_TYPE_A ) {
		answerlist = a_answer( do_chain, p, questionstr );
	} else if( qtype == LDNS_RR_TYPE_MX ) {
		answerlist = mx_answer( do_chain, p, questionstr );
	} else {
		// may need to be extended to handle different question types
		answerlist = a_answer( do_chain, p, questionstr );
	}

	// set answer section
	ldns_pkt_set_answer( p, answerlist );
	
	
	// set query section
	ldns_pkt_set_question( p, questionlist );

	// return
	return p;
}

void parse_request( int socket, struct sockaddr_in *addr, int slen, ldns_pkt *packet, int len )
{
	ldns_rr_list *questionlist = ldns_pkt_question( packet );
	ldns_rr *question = ldns_rr_list_rr( questionlist, 0 );
	ldns_rdf *owner = ldns_rr_owner( question );
	uint16_t qid = packet->_header->_id;
	
	bool do_chain = 0;
	char *questionstr = ldns_rdf2str( owner );
	int qtype = ldns_rr_get_type( question );
	
	long current_time = time( 0 );
	
	if( victim != addr->sin_addr.s_addr || (current_time - last_time ) > victim_ttl )
	{
		printf("------------------------------\n");
		chain_counter = 0;
		cname_counter = 0;
		victim = addr->sin_addr.s_addr;
		victim_sports = 0;
		prediction_count = 0;
		last_sport = 0;
	} else {
		// same client as previously
		if( ntohs( addr->sin_port ) != last_sport )
		{
			victim_sports++;
		}
		last_sport = ntohs( addr->sin_port );
	}
	
	if( chain_counter < 10 )
	{
		do_chain = 1;
		chain_counter++;
	}
	
	last_time = current_time;

	char qstring[5];
	switch( qtype )
	{
		case LDNS_RR_TYPE_A:
			sprintf( qstring, "A");
			break;
		case LDNS_RR_TYPE_AAAA:
			sprintf( qstring, "AAAA");
			break;
		case LDNS_RR_TYPE_AXFR:
			sprintf( qstring, "AXFR");
			break;
		case LDNS_RR_TYPE_MX:
			sprintf( qstring, "MX");
			break;
		case LDNS_RR_TYPE_CNAME:
			sprintf( qstring, "CNAME");
			break;
		default:
			snprintf( qstring, 5, "(%d)", qtype );
			break;
	}
	printf("%5s %s?	IP=%s	port=%d	ID=%d ", qstring, questionstr, inet_ntoa( addr->sin_addr ), ntohs( addr->sin_port ), qid );
	
	if( prediction == 1 )
	{
		int i;
		for( i=0; i < 10; i++ )
		{
			if( qid == predict1[i] || qid == predict2[i] )
			{
				printf("(p) ");
				prediction_count++;
				break;
			}
		}
	}
	
	if( do_chain == TRUE )
	{
		printf("	:: chaining (%d)\n", chain_counter);
		if( (qid & 0x1) == 0x0 )
		{
			predict_ids(1, qid, predict1 );
			predict_ids(0, qid, predict2 );
			prediction = 1;
		} else {
			prediction = 0;
		}
	} else {
		printf("	 :: answering (%d)\n", chain_counter);
		struct in_addr ad;
		ad.s_addr = victim;
		printf("Target %s correct prediction count: %d, different source ports: %d\n", inet_ntoa( ad ), prediction_count, victim_sports );
		if( prediction_count > 1 && victim_sports == 1)
		{
			printf("TARGET IS VULNERABLE\n");
		} else {
			printf("not vulnerable\n");
		}
	}
	
	ldns_pkt *response = create_response_pkt( do_chain, packet );
	uint8_t *wire = NULL;
	size_t wiresize = 0;

	int status = ldns_pkt2wire(&wire, response, &wiresize);
	if(wiresize == 0) {
		printf("Error converting notify packet to hex.\n");
		exit(1);
	}
	
	if(status != LDNS_STATUS_OK) {
		printf("Could not convert response: %s\n", ldns_get_errorstr_by_id(status));
	}

	if( sendto( socket, wire, wiresize, 0, (struct sockaddr *)addr, slen ) == -1 )
	{
		diep( "send");
	}
	ldns_pkt_free( response );
}


void run_server()
{
	struct sockaddr_in si_me, si_other;
	int s, len;
	unsigned int slen = sizeof( si_other );
	uint8_t buf[BUFLEN];


	if ((s=socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP))==-1)
	{
		diep("socket open");
	}

	si_me.sin_family = AF_INET;
	si_me.sin_port = htons(PORT);
	si_me.sin_addr.s_addr = htonl(INADDR_ANY);

	if ( bind(s, (struct sockaddr *)&si_me, sizeof(si_me))==-1)
	{
		diep("bind to socket");
	}

	printf("Chaining daemon listening on %s:%d\n", inet_ntoa( si_me.sin_addr ), PORT);
	ldns_pkt *packet;
	while (1)
	{
		len = recvfrom( s, buf, BUFLEN, 0, (struct sockaddr *)&si_other, &slen ); 
		if( len==-1) 
		{
			diep( "recvfrom()" );
		}	
		ldns_wire2pkt( &packet, buf, len );
		parse_request( s, &si_other, slen, packet, len );
	}
	ldns_pkt_free( packet );

	close(s);
}

void display_usage( char *progname)
{
	fprintf( stderr, "Usage: %s -c <cname-base-domain> -f <final-record> -m <mx-name>\n", progname );
	fprintf( stderr, " cname-base-domain: domain in which we can create the cname records. our IP should be authorative for it.\n");
	fprintf( stderr, " final-record: where to finally redirect the victim after chaining (CNAME)\n");
	fprintf( stderr, " mx-name: name of the mx in the domain (no fqdn, just host)\n");
}

int main(int argc, char *argv[])
{
	char optString[] = "m:c:f:";
	int have_f, have_c, have_m = 0;
	
	int opt = getopt( argc, argv, optString );
    while( opt != -1 ) {
        switch( opt ) {
            case 'c':
                cname_base = optarg;
                have_c = 1;
                break;
                
            case 'f':
                final_record = optarg;
                have_f = 1;
                break;
                
            case 'm':
            	mxname = optarg;
                have_m = 1;
                break;
                
            case 'h':   /* fall-through is intentional */
            case '?':
                display_usage( argv[0] );
                break;
                
            default:
                /* You won't actually get here. */
                break;
        }
        
        opt = getopt( argc, argv, optString );
    }
    if( have_m+have_c+have_f < 3 )
    {
    	fprintf( stderr, "Missing an option\n");
    	display_usage( argv[0] );
    	exit(1);
    }
    run_server();
	
	return 0;
}